package transfer

import (
	"Infinichat_mixed/entity"
	"bytes"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"sync"
)

/**
* @author aFun
* @data 2025/4/18 18:20
* @email: afunaaa222@163.com / 1030935903
* @address: Guangdong University of Technology, China
 */

const (
	maxChunkSize = 4096
	maxPkgSize   = 100 << 20
	headerSize   = 4
)

type Transfer struct {
	Conn net.Conn
}

// 使用分层缓冲池
var (
	headerPool = sync.Pool{
		New: func() interface{} {
			return make([]byte, headerSize)
		},
	}
	chunkPool = sync.Pool{
		New: func() interface{} {
			return make([]byte, maxChunkSize)
		},
	}
)

func (t *Transfer) ReadPkg() (mes *entity.Message, err error) {
	lenBuf := headerPool.Get().([]byte)
	defer headerPool.Put(lenBuf[:headerSize]) // 确保放回时是原始大小

	if _, err := io.ReadFull(t.Conn, lenBuf[:headerSize]); err != nil {
		return &entity.Message{
			Status:  false,
			Message: "读取消息头失败",
		}, fmt.Errorf("read header failed: %w", err)
	}

	pkgLen := binary.BigEndian.Uint32(lenBuf[:headerSize])
	if pkgLen == 0 || pkgLen > maxPkgSize {
		return &entity.Message{
			Status:  false,
			Message: fmt.Sprintf("消息过大(%d > %d)", pkgLen, maxPkgSize),
		}, fmt.Errorf("包大小 %d 超出限制", pkgLen)
	}
	bodyBuf := make([]byte, pkgLen)
	if pkgLen > maxChunkSize {
		// 分片读取
		chunk := chunkPool.Get().([]byte)
		defer chunkPool.Put(chunk[:maxChunkSize])
		var buf bytes.Buffer
		for offset := 0; offset < int(pkgLen); {
			readSize := min(int(pkgLen)-offset, maxChunkSize)
			if _, err := io.ReadFull(t.Conn, chunk[:readSize]); err != nil {
				return &entity.Message{
					Status:  false,
					Message: "读取消息体失败",
				}, fmt.Errorf("read body chunk failed: %w", err)
			}
			//copy(bodyBuf[offset:], chunk[:readSize])
			buf.Write(chunk[:readSize])
			offset += readSize
		}
		bodyBuf = buf.Bytes()
	} else {
		if _, err := io.ReadFull(t.Conn, bodyBuf); err != nil {
			return &entity.Message{
				Status:  false,
				Message: "读取消息体失败",
			}, fmt.Errorf("read body failed: %w", err)
		}
	}

	var msg entity.Message

	if err := json.Unmarshal(bodyBuf, &msg); err != nil {
		fmt.Printf("bodybuf: %s\n", bodyBuf)
		return &entity.Message{
			Status:  false,
			Message: "消息解析失败",
		}, fmt.Errorf("unmarshal failed: %w", err)
	}

	return &msg, nil
}

func (t *Transfer) WriterPkg(data []byte) (err error) {
	if len(data) > maxPkgSize {
		return fmt.Errorf("消息过大(%d > %d)", len(data), maxPkgSize)
	}

	lenBuf := headerPool.Get().([]byte)
	defer headerPool.Put(lenBuf[:headerSize])

	binary.BigEndian.PutUint32(lenBuf[:headerSize], uint32(len(data)))
	if _, err := t.Conn.Write(lenBuf[:headerSize]); err != nil {
		return fmt.Errorf("写头失败: %w", err)
	}

	if len(data) > maxChunkSize {
		chunk := chunkPool.Get().([]byte)
		defer chunkPool.Put(chunk[:maxChunkSize])

		for offset := 0; offset < len(data); {
			writeSize := min(len(data)-offset, maxChunkSize)
			copy(chunk, data[offset:offset+writeSize])
			if _, err := t.Conn.Write(chunk[:writeSize]); err != nil {
				return fmt.Errorf("写体失败: %w", err)
			}
			offset += writeSize
		}
	} else {
		if _, err := t.Conn.Write(data); err != nil {
			return fmt.Errorf("写体失败: %w", err)
		}
	}

	return nil
}
